package com.jcxt.sys.common.code_generate;

import com.baomidou.mybatisplus.generator.AutoGenerator;
import com.baomidou.mybatisplus.generator.InjectionConfig;
import com.baomidou.mybatisplus.generator.config.*;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.config.rules.NamingStrategy;
import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.ToString;
import lombok.val;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.config.YamlMapFactoryBean;
import org.springframework.core.io.ClassPathResource;

import java.nio.file.Paths;
import java.util.Arrays;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import static org.apache.commons.lang3.SystemUtils.USER_DIR;


/**
 * @author 邱理
 * @description
 * @date 创建于 2019/4/16
 */
public class MybatisCodeGenerator {

    public static String getSrcDir(String projectDir) {
        return Paths.get(USER_DIR, projectDir == null ? "":projectDir, "src/main/java").toString();
    }

    @Data
    @EqualsAndHashCode(callSuper = true)
    public static class CustomFileOutConfig extends FileOutConfig {

        private String outputDir;

        private PackageConfig packageConfig;

        private CustomGlobalConfig customGlobalConfig;

        public CustomFileOutConfig() {
            this.setTemplatePath("mybatis/templates/mapper.xml.ftl");
        }

        @Override
        public String outputFile(TableInfo tableInfo) {
            String path = Paths.get(USER_DIR,
                    customGlobalConfig.getProjectDir() == null ? "" : customGlobalConfig.getProjectDir(),
                    outputDir,
                    packageConfig.getModuleName(),
                    tableInfo.getEntityName() + "Mapper.xml").toString();

            return path;
        }
    }

    @Data
    @EqualsAndHashCode(callSuper = true)
    @ToString(callSuper = true)
    public static class CustomGlobalConfig extends GlobalConfig {

        private String projectDir;

        @Override
        public String getOutputDir() {
            //if (projectDir == null) return super.getOutputDir();
            return getSrcDir(projectDir);
        }
    }

    @Data
    @EqualsAndHashCode(callSuper = true)
    @ToString(callSuper = true)
    public static class CustomDataSourceConfig extends DataSourceConfig {
        private String springConfigProfile;
        private ClassLoader classLoader;

        public void setConfigBySpringConfigProfile() {
            //if (springConfigProfile != null) {
                YamlMapFactoryBean yamlConfigBean = new YamlMapFactoryBean();
                String suffix = "";
                if (!StringUtils.isBlank(springConfigProfile)) {
                    suffix = "-" + springConfigProfile;
                }
                yamlConfigBean.setResources(new ClassPathResource("application" + suffix + ".yml", classLoader));
                Map<String, Object> config = yamlConfigBean.getObject();
                val datasource = (Map<String, Object>) ((Map<String, Object>) ((Map<String, Object>) config.get("spring")).get("datasource")).get("druid");
                this.setUrl(datasource.get("url").toString())
                        .setDriverName(datasource.get("driver-class-name").toString())
                        .setUsername(datasource.get("username").toString())
                        .setPassword(datasource.get("password").toString());
            }
       // }
    }

    @Data
    @EqualsAndHashCode(callSuper = true)
    @ToString(callSuper = true)
    public static class CustomStrategyConfig extends StrategyConfig {

        private String includes;

        @Override
        public Set<String> getInclude() {
            if (StringUtils.isBlank(includes)) return super.getInclude();
            return Arrays.stream(includes.split(",")).collect(Collectors.toSet());
        }
    }

    public static class CustomInjectConfig extends InjectionConfig {

        @Override
        public void initMap() {

        }
    }


    public static AutoGenerator buildAutoGenerator(ClassLoader cl, String propertiesFile) {

        PropertyFileInjector fileInjector = new PropertyFileInjector(cl, propertiesFile);
        CustomGlobalConfig globalConfig = fileInjector.inject(CustomGlobalConfig.class, "global");
        CustomDataSourceConfig dataSourceConfig = fileInjector.inject(CustomDataSourceConfig.class, "datasource");
        dataSourceConfig.setClassLoader(cl);
        dataSourceConfig.setConfigBySpringConfigProfile();
        PackageConfig packageConfig = fileInjector.inject(PackageConfig.class, "package");
        packageConfig.setEntity("domain");
        TemplateConfig templateConfig = fileInjector.inject(TemplateConfig.class, "template");
        templateConfig.setXml(null);
        CustomStrategyConfig strategyConfig = fileInjector.inject(CustomStrategyConfig.class, "strategy");
        strategyConfig.setNaming(NamingStrategy.underline_to_camel);
        strategyConfig.setColumnNaming(NamingStrategy.underline_to_camel);

        InjectionConfig cfg = new InjectionConfig() {
            @Override
            public void initMap() {

            }
        };

        CustomFileOutConfig fileOutConfig = fileInjector.inject(CustomFileOutConfig.class, "xml");
        fileOutConfig.setCustomGlobalConfig(globalConfig);
        fileOutConfig.setPackageConfig(packageConfig);

        cfg.setFileOutConfigList(Arrays.asList(fileOutConfig));

        AutoGenerator ag = new AutoGenerator();

        ag.setGlobalConfig(globalConfig)
                .setCfg(cfg)
                .setDataSource(dataSourceConfig)
                .setPackageInfo(packageConfig)
                .setTemplate(templateConfig)
                .setStrategy(strategyConfig)
                .setTemplateEngine(new FreemarkerTemplateEngine());

        return ag;
    }

}
